package com.treeroot.treenet.protocol.abs.tcp.session;

import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.SocketChannel;
import java.util.Queue;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * Description: Description: TCP写缓冲队列，支持高并发写入，单线程写出
 * Create By TreeRoot
 * Create Time: 2025/06/02 10:48
 * Class Name: WriteBufferQueue
 *
 * @author TreeRoot
 */
@Slf4j
public class WriteBufferQueue {

    private final SocketChannel channel;
    private final Queue<ByteBuffer> writeQueue = new ConcurrentLinkedQueue<>();
    private final AtomicBoolean writing = new AtomicBoolean(false);

    public WriteBufferQueue(SocketChannel channel) {
        this.channel = channel;
    }

    /**
     * 高并发安全写入队列
     */
    public void offer(ByteBuffer data) {
        writeQueue.offer(data);
        tryWrite();
    }

    /**
     * 尝试写出数据（串行执行，防止多线程同时写 channel）
     */
    private void tryWrite() {
        if (!writing.compareAndSet(false, true)) {
            // 正在写，避免重复触发
            return;
        }

        try {
            ByteBuffer buffer;
            while ((buffer = writeQueue.peek()) != null) {
                channel.write(buffer);
                if (buffer.hasRemaining()) {
                    // Socket 写缓冲区满了，下次再写
                    break;
                }
                // 已写完，移除
                writeQueue.poll();
            }
        } catch (IOException e) {
            log.error("[WriteBufferQueue] 写入异常，关闭连接", e);
            try {
                channel.close();
            } catch (IOException ex) {
                log.warn("[WriteBufferQueue] 关闭异常", ex);
            }
        } finally {
            writing.set(false);

            // 如果写完后又有数据进来，重新触发写
            if (!writeQueue.isEmpty()) {
                tryWrite();
            }
        }
    }
}